import os
import umap
import time
import random
import anndata
import argparse
import numpy as np
import pandas as pd
import scanpy as sc
import editdistance
import colorsys as cs
import seaborn as sns
import memory_profiler
import scipy.sparse as sp
from collections import Counter
import matplotlib.pyplot as plt
from scipy.spatial import procrustes
from sklearn.decomposition import PCA
from scipy.spatial.distance import cdist
from matplotlib.cm import ScalarMappable
from sklearn.neighbors import NearestNeighbors
from scipy.linalg import orthogonal_procrustes
from matplotlib.colors import LinearSegmentedColormap, Normalize
def make_colormap(seq):
seq = [(None,) * 3, 0.0] + list(seq) + [1.0, (None,) * 3]
cdict = {'red': [], 'green': [], 'blue': []}
for i, item in enumerate(seq):
if isinstance(item, float):
r1, g1, b1 = seq[i - 1]
r2, g2, b2 = seq[i + 1]
cdict['red'].append([item, r1, r2])
cdict['green'].append([item, g1, g2])
cdict['blue'].append([item, b1, b2])
return LinearSegmentedColormap('CustomMap', cdict)
def build_6mer_dist(bc_list):
start_km = {}
mid_km = {}
end_km = {}
for bc in bc_list:
start_km.setdefault(bc[:6] , []).append(bc)
mid_km.setdefault(bc[4:10], []).append(bc)
end_km.setdefault(bc[-6:] , []).append(bc)
return start_km,mid_km,end_km
def barcode_matching(bc_pos_dict,spatial_bc_list,max_dist=1):
bc_matching_dict = {}
def get_sel_bc(bc):
res = []
if bc[:6] in start_km:
res += start_km[bc[:6]]
if bc[-6:] in end_km:
res += end_km[bc[-6:]]
if bc[4:10] in mid_km:
res += mid_km[bc[4:10]]
return set(res)
exact_match = 0
fuzzy_match = 0
bc_ref_list = list(bc_pos_dict.keys())
start_km,mid_km,end_km = build_6mer_dist(bc_ref_list)
for bc in spatial_bc_list:
bc_old = bc
#bc = "".join([bc[it] for it in [1,2,3,4,5,6,7,9,10,11,12,13] ])
if bc in bc_pos_dict:
exact_match += 1
bc_matching_dict[bc_old] = bc
else:
sel_bc = get_sel_bc(bc)
#sel_bc = bc_ref_list
if len(sel_bc)>0:
fz = [(it, editdistance.eval(it, bc)) for it in sel_bc]
fz = [it for it in fz if it[1]<=max_dist]
fz.sort(key=lambda x:x[1])
if len(fz)==0:
continue
## if there are two barcodes with the same edit distance, choose the one with higher error rate in the last base
if len(fz)>1 and fz[0][1]==fz[1][1]:
if editdistance.eval(fz[0][0][:-1], bc[-1])>editdistance.eval(fz[1][0][:-1], bc[-1]): # higher error rate in the last base of the barcode
fuzzy_match += 1
bc_matching_dict[bc_old] = fz[1][0]
elif editdistance.eval(fz[0][0][:-1], bc[-1])<editdistance.eval(fz[1][0][:-1], bc[-1]):
fuzzy_match += 1
bc_matching_dict[bc_old] = fz[0][0]
else:
fuzzy_match += 1
bc_matching_dict[bc_old] = fz[0][0]
return bc_matching_dict,exact_match,fuzzy_match
def blind_cnt_distribution(bead_all, bead_type):
plt.figure(figsize=(8,6))
sns.histplot(np.log10(bead_all['total_cnt']), bins=50)
plt.xlabel('log10(total count)', fontsize=12)
plt.ylabel('number of '+bead_type, fontsize=12)
plt.title('blind '+bead_type+' total count distribution ({}), median={}'.format(len(bead_all), bead_all['total_cnt'].median()), fontsize=16, pad=20)
plt.tick_params(axis='both', which='major', labelsize=10)
plt.gca().spines['top'].set_visible(False)
plt.gca().spines['right'].set_visible(False)
plt.savefig(os.path.join(out_dir,bead_type+'_blind_cnt_distribution.png'),dpi=300)
plt.close()
def blind_cover_bc_distribution(bead_cover_bc, bead_type):
plt.figure(figsize=(8,6))
sns.histplot(bead_cover_bc['cnt'], bins=50)
plt.xlabel('bead covered', fontsize=12)
plt.ylabel('number of '+bead_type, fontsize=12)
plt.title(bead_type+' bead covered distribution ({}), median={}'.format(len(bead_cover_bc), bead_cover_bc['cnt'].median()), fontsize=16, pad=20)
plt.tick_params(axis='both', which='major', labelsize=10)
plt.gca().spines['top'].set_visible(False)
plt.gca().spines['right'].set_visible(False)
plt.savefig(os.path.join(out_dir,bead_type+'_blind_cover_bc_distribution.png'),dpi=300)
plt.close()
# generate spase matrix from matching, with selection on anchor or target
def get_matrix(match_df,min_a_cnt,max_a_cnt, min_t_cnt,max_t_cnt, anchor, target):
a_all = match_df.groupby(anchor)['cnt'].sum().reset_index(name='total_cnt')
a_sel = a_all.loc[(a_all['total_cnt']>min_a_cnt) & (a_all['total_cnt']<max_a_cnt),]
t_all = match_df.groupby(target)['cnt'].sum().reset_index(name='total_cnt')
t_sel = t_all.loc[(t_all['total_cnt']>min_t_cnt) & (t_all['total_cnt']<max_t_cnt),]
match_df = match_df[(match_df[anchor].isin(a_sel[anchor])) & (match_df[target].isin(t_sel[target]))]
a_list = match_df.groupby(anchor)['cnt'].sum().reset_index(name='total_cnt')
t_list = match_df.groupby(target)['cnt'].sum().reset_index(name='total_cnt')
print('a: {}'.format(len(a_list)))
print('t: {}'.format(len(t_list)))
a_dict = dict()
t_dict = dict()
for i in range(len(a_list)):
a_dict[a_list.iloc[i,0]] = i
for j in range(len(t_list)):
t_dict[t_list.iloc[j,0]] = j
a_coo = []
t_coo = []
[a_coo.append(a_dict[a]) for a in match_df[anchor]]
[t_coo.append(t_dict[t]) for t in match_df[target]]
counts_coo = sp.coo_matrix((match_df['cnt'], (a_coo, t_coo)))
# counts = counts_coo.tocsr().toarray()
counts = counts_coo.tocsr()
return counts, a_list, t_list
def get_col(coords):
plot_df = coords.copy()
b_list = [1] * len(plot_df)
r_list = (plot_df.xcoord.values - plot_df.xcoord.min()) / (
plot_df.xcoord.max() - plot_df.xcoord.min())
g_list = (plot_df.ycoord.values - plot_df.ycoord.min()) / (
plot_df.ycoord.max() - plot_df.ycoord.min())
# my_h = plot_df['xcoord']
# my_h = (my_h - np.min(my_h))
# my_h = 0.2+my_h * (0.8 / np.max(my_h))
# my_s = np.ones(len(my_h)) * 0.5
# my_l = plot_df['ycoord']
# my_l = (my_l - np.min(my_l))
# my_l = 0.2 + my_l * (0.7 / np.max(my_l))
# hls_color = np.column_stack((my_h, my_l, my_s))
# rgb_color = [cs.hls_to_rgb(p[0], p[1], p[2]) for p in hls_color]
rgb_color = [(r_list[i], g_list[i], b_list[i]) for i in range(len(plot_df))]
return rgb_color
def get_truth(wlist,pos_df):
w_truth = pos_df[pos_df['barcode'].isin(wlist)]
w_truth = w_truth.sort_values(by=['barcode'])
w_col = get_col(w_truth)
return w_truth, w_col
def get_pa(pos_truth,pos_recon):
mtx1, mtx2, disparity = procrustes(pos_truth[['xcoord','ycoord']], pos_recon[['xcoord','ycoord']])
pos_truth_translate = pos_truth.copy()
pos_truth_translate['xcoord'] = pos_truth_translate['xcoord'] - np.mean(pos_truth_translate['xcoord'])
pos_truth_translate['ycoord'] = pos_truth_translate['ycoord'] - np.mean(pos_truth_translate['ycoord'])
scaling = np.sqrt(np.trace(np.dot(pos_truth_translate[['xcoord','ycoord']], pos_truth_translate[['xcoord','ycoord']].T)))
scaling_2 = np.sqrt(np.trace(np.dot(pos_truth_translate[['xcoord','ycoord']], pos_truth_translate[['xcoord','ycoord']].T)))/np.sqrt(np.trace(np.dot(mtx2, mtx2.T)))
mtx1_scaled = mtx1*scaling*0.72
mtx2_scaled = mtx2*scaling_2*0.72
dist = pos_truth_translate.copy()
dist['x'] = mtx1_scaled[:,0] - mtx2_scaled[:,0]
dist['y'] = mtx1_scaled[:,1] - mtx2_scaled[:,1]
dist['r'] = np.sqrt(dist['x']**2 + dist['y']**2)
return disparity, mtx1_scaled, mtx2_scaled, dist
def rigid_transfrom(pos, a_recon):
bc_pos_dict = Counter(pos['barcode'])
a_bc_matching_dict,_,_ = barcode_matching(bc_pos_dict, a_recon[anchor].values, max_dist=1)
a_truth, a_col = get_truth(list(a_bc_matching_dict.values()), pos)
a_umap_matched = a_recon[a_recon[anchor].isin(a_bc_matching_dict.keys())]
a_umap_matched = a_umap_matched.copy()
a_umap_matched.loc[:,'barcode'] = a_umap_matched[anchor].map(a_bc_matching_dict)
a_umap_matched = a_umap_matched[['xcoord','ycoord','barcode']]
a_umap_matched = a_umap_matched.groupby('barcode').mean().reset_index()
a_umap_matched = a_umap_matched.sort_values(by=['barcode'])
mtx1, mtx2, disparity = procrustes(a_truth[['xcoord','ycoord']], a_umap_matched[['xcoord','ycoord']])
a_matched_translate = a_umap_matched.copy()
a_matched_translate['xcoord'] = a_matched_translate['xcoord'] - np.mean(a_matched_translate['xcoord'])
a_matched_translate['ycoord'] = a_matched_translate['ycoord'] - np.mean(a_matched_translate['ycoord'])
R, s = orthogonal_procrustes(a_matched_translate[['xcoord','ycoord']], mtx2)
a_transformed = s * np.dot(a_recon[['xcoord','ycoord']], R)
a_scale = a_recon.copy()
a_scale['xcoord'] = a_transformed[:,0]/(np.max(a_transformed[:,0])-np.min(a_transformed[:,0]))*(np.max(pos.xcoord)-np.min(pos.xcoord))*0.72
a_scale['ycoord'] = a_transformed[:,1]/(np.max(a_transformed[:,1])-np.min(a_transformed[:,1]))*(np.max(pos.ycoord)-np.min(pos.ycoord))*0.72
a_scale['xcoord'] = a_scale['xcoord'] - np.mean(a_scale['xcoord'])
a_scale['ycoord'] = a_scale['ycoord'] - np.mean(a_scale['ycoord'])
return a_scale
